/**
 * Wuhan University Library for Cryptography (WULC) is an ongoing and portable 
 * cryptographic toolkit developed by Cryptography and Blockchain Technology
 * Laboratory · Wuhan University. 
 * This file is part of WULC. 
 * @copyright Copyright (c) 2018-2023 WULC authors
 * 
 * @file pemks_function.java
 * @brief to descirbe pemks_function
 * 
 * @author Qiao Hongyi (974446987@qq.com)
 * @version 0.1.1
 * @date 2023-01-16
 * 
 * @ingroup {group}
 */
package pemks;

import com.hankcs.hanlp.HanLP;
import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.Field;
import it.unisa.dia.gas.jpbc.Pairing;
import it.unisa.dia.gas.plaf.jpbc.pairing.PairingFactory;
import org.zz.gmhelper.SM4Util;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.*;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.*;


public class pemks_function {
    private static final char[] HEX_ARRAY = "0123456789ABCDEF".toCharArray();
    static int SM4_BLOCK_SIZE = 16;
    static boolean SM4algorithm = true;//1
    static byte[] sm4_cipherText = null;
    static byte[] sm4_decryptedData = null;
    static int l1 = 5;

    public static Pairing bp = PairingFactory.getPairing("a.properties");

    public  static  Field G1 = bp.getG1();
    public  static  Field G2 = bp.getG2();
    public  static  Field GT = bp.getGT();
    public  static  Field Zq = bp.getZr();



    public static Element g1 = G1.newRandomElement().getImmutable();
    public static Element g2 = G2.newRandomElement().getImmutable();

    public static Map<String, Element> key_map = KeyGen(Zq,g1,g2);
    public static Element g1_e_alpha = key_map.get("sk"); //用户私钥
    public static Element g2_e_alpha = key_map.get("pk_h");
    public static void main(String[] args) throws NoSuchProviderException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, NoSuchPaddingException, InvalidKeyException {
        // 一、基于特定椭圆曲线类型生成Pairing实例
        // 1.从文件导入椭圆曲线参数
        Pairing bp = PairingFactory.getPairing("a.properties");

        Field G1 = bp.getG1();
        Field G2 = bp.getG2();
        Field GT = bp.getGT();
        Field Zq = bp.getZr();

        Element g1 = G1.newRandomElement().getImmutable();
        Element g2 = G2.newRandomElement().getImmutable();

/**            B. Key Generation           **/
        Map<String, Element> key_map = KeyGen(Zq,g1,g2);
        Element g1_e_alpha = key_map.get("sk"); //用户私钥
        Element g2_e_alpha = key_map.get("pk_h"); //h，用户公钥为{g1, g2, h}.

/**            C. Encryption(Encryption comprises file encryption and index encryption.)**/
        System.out.println("DO请输入要加密的明文文件:");      //C:\\Users\\97444\\Desktop\\a.txt
        Scanner sc=new Scanner(System.in);
        String plainfile = "";
        if(sc.hasNextLine()){
            plainfile = sc.nextLine();
        }
        Map<String, String> ct_map = encryption(GT, bp, Zq, g1, g2, g2_e_alpha, plainfile);
        String secret_kf = ct_map.get("kf"); //文件私钥
        String CT = ct_map.get("final_ct");

/**            D. Trapdoor Generation         **/
        System.out.println("User请输入要查询的关键词集合（每个关键词之间用逗号分隔）:");//算法,处理,通信
        Scanner sc2=new Scanner(System.in);
        String testwords = "";
        if(sc2.hasNextLine()){
            testwords = sc2.nextLine();
        }
        sc.close();
        sc2.close();

        String Tw = trapdoor(testwords, Zq, g1, g1_e_alpha, bp);

/**             E. Search          **/
        Map<String, String> search_map = Search(bp, G1, G2, GT, Zq, CT, Tw);

        if(search_map.get("flag") == "true"){
            System.out.println("查询成功！解密得到的明文是：");
            String dp = "D:\\Code\\job\\order\\cipher-cloud\\c.txt";//解密后文件
            file_decryption(secret_kf, search_map.get("filecontent"), dp);

        }else if(search_map.get("flag") == "false"){
            System.out.println("查询失败！");
        }

    }

    public static String file_decryption(String secret_kf, String filecontent, String dp) throws BadPaddingException, NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, NoSuchProviderException, InvalidKeyException {
        /**            解密文件部分           **/
        byte[] bs = new byte[SM4_BLOCK_SIZE];
        System.arraycopy(hexStringToBytes(secret_kf), 0, bs, 0, SM4_BLOCK_SIZE);
        sm4_decryptedData = SM4Util.decrypt_ECB_Padding(bs, hexStringToBytes(filecontent));
        String b = new String(sm4_decryptedData);
        System.out.println(b);
        writeFileContent(dp, b);
        return b;
    }

    public static Map<String, String> Search(Pairing bp, Field G1, Field G2, Field GT, Field Zq, String CT, String Tw) {
        /**      e(T2, C2) = C1 · e(T1, ΠCi^Ti)       **/
        //分割
        String[] CT_split = CT.split(",");
        String[] Tw_split = Tw.split(",");

        String C1 = CT_split[0];
        String C2 = CT_split[1];
        List<String> Ci_list = new ArrayList<String>();
        for (int i = 0; i <= l1; i++) {
            Ci_list.add(CT_split[i+2]);
        }
        String CF = CT_split[CT_split.length - 1];

        String T1 = Tw_split[0];
        String T2 = Tw_split[1];
        List<String> Ti_list = new ArrayList<String>();
        for (int i = 2; i <= Tw_split.length - 1; i++) {
            Ti_list.add(Tw_split[i]);
        }

        //计算
        Element C1_prime = GT.newRandomElement().setToZero();
        Element C2_prime = G2.newRandomElement().setToZero();
        Element T1_prime = G1.newRandomElement().setToZero();
        Element T2_prime = G1.newRandomElement().setToZero();
        C1_prime.setFromBytes(hexStringToBytes(C1));
        C2_prime.setFromBytes(hexStringToBytes(C2));
        T1_prime.setFromBytes(hexStringToBytes(T1));
        T2_prime.setFromBytes(hexStringToBytes(T2));

        Element T2_C2 = bp.pairing(T2_prime.duplicate(), C2_prime.duplicate());     //e(T2, C2)

        Element gc = G2.newRandomElement().setToZero();
        Element gt = Zq.newRandomElement().setToZero();
        gc.setFromBytes(hexStringToBytes(Ci_list.get(0)));
        gt.setFromBytes(hexStringToBytes(Ti_list.get(0)));

        Element Ci_e_Ti = gc.duplicate().powZn(gt.duplicate());   //C0^T0

        for (int i = 1; i <= l1; i++) {
            gc.setFromBytes(hexStringToBytes(Ci_list.get(i)));
            gt.setFromBytes(hexStringToBytes(Ti_list.get(i)));
            Ci_e_Ti = Ci_e_Ti.duplicate().mul(gc.duplicate().powZn(gt.duplicate()));
        }

        Element T1_CiTi = bp.pairing(T1_prime.duplicate(), Ci_e_Ti.duplicate());   //e(T1, ΠCi^Ti)
        Element C1_m_T1_CiTi = C1_prime.duplicate().mul(T1_CiTi.duplicate());      //C1 · e(T1, ΠCi^Ti)

        Map<String, String> map = new HashMap<String, String>();
        if (T2_C2.isEqual(C1_m_T1_CiTi)) {
            map.put("flag", "true");
            map.put("filecontent", CF);
        } else {
            map.put("flag", "false");
            map.put("filecontent", "");
        }
        return map;
    }

    public static String trapdoor(String testwords, Field Zq, Element g1, Element g1_e_alpha, Pairing bp) {
        String[] test_keyword_Arr = testwords.split(",");
        List<String> test_keywordList = new ArrayList<String>();  //用户选择的关键词集合
        for(int i = 0; i<test_keyword_Arr.length; i++){
            test_keywordList.add(test_keyword_Arr[i]);
        }

        int l2 = test_keywordList.size();
/**            1)user choose W0 = {w1, w2, ..., wl2}, l2<=l1(5)
 *             selects two random numbers r, β ∈ Zq∗ and computes T1 = g1^(r/β), T2 = (g1^α)(g1^l2r)
 **/
        Element r = Zq.newRandomElement().getImmutable();
        Element β = Zq.newRandomElement().getImmutable();

        Element T1_prime = g1.duplicate().powZn(r.duplicate().div(β.duplicate()));         //g1^(r/β)
        Element g1_e_l2r = g1.duplicate().powZn(r.duplicate().mul(l2));
        Element T2_prime = g1_e_alpha.duplicate().mul(g1_e_l2r.duplicate());               //(g1^α)(g1^l2r)

        String T1 = hexBytesToString(T1_prime.duplicate().toBytes());  //Tw---T1
        String T2 = hexBytesToString(T2_prime.duplicate().toBytes());  //Tw---T2

/**            2)calculates Ti          **/
        List<BigInteger> test_bHw_list = new ArrayList<BigInteger>();
        for (Iterator<String> it = test_keywordList.iterator(); it.hasNext(); ) {
            byte[] tmp = hashFromStringToZp(it.next(), bp).toBytes();  //H'(w1), H'(w2), ..., H'(wl2)
            String new_Hw = hexBytesToString(tmp);
            BigInteger bigInt = new BigInteger(new_Hw, 16);
            test_bHw_list.add(bigInt);  //H'(w1), H'(w2), ..., H'(wl2)
        }

        String Tw = T1 + "," + T2;
        for (int i = 0; i <= l1; i++) {
            Element Ti = Zq.newRandomElement();
            Ti.setToZero();
            for (int j = 0; j < l2; j++) {
                Ti = Ti.duplicate().add(β.duplicate().mul(test_bHw_list.get(j).pow(i)));
            }
            Tw = Tw + "," + hexBytesToString(Ti.duplicate().toBytes());
        }
        System.out.println(Tw);
        return Tw;
    }

    public static Map<String, Element> KeyGen(Field zq, Element g1, Element g2) {
        Element alpha = zq.newRandomElement().getImmutable();

        Element g1_e_alpha = g1.duplicate().powZn(alpha.duplicate()); //用户私钥
        Element g2_e_alpha = g2.duplicate().powZn(alpha.duplicate()); //h，用户公钥为{g1, g2, h}.

        Map<String, Element> map = new HashMap<String, Element>();
        map.put("sk", g1_e_alpha);
        map.put("pk_h", g2_e_alpha);
        return map;
    }

    public static Map<String, String> encryption(Field GT, Pairing bp, Field zq, Element g1, Element g2, Element g2_e_alpha, String plainfile) throws NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, NoSuchProviderException, InvalidKeyException {
        Map<String, String> fe_map = file_encryption(GT,bp,plainfile);
        String CF = fe_map.get("ciph");
        String kf = fe_map.get("KeyF");



        Map<String, String> ct_map = index_encryption(bp,zq,g1,g2,g2_e_alpha,CF,plainfile);
        String final_ct = ct_map.get("CT");
        String keywords_string = ct_map.get("keywords");

        Map<String, String> map = new HashMap<String, String>();
        map.put("kf", kf);
        map.put("final_ct", final_ct);
        map.put("keywords", keywords_string);
        return map;
    }

    public static Map<String, String> index_encryption(Pairing bp, Field Zq, Element g1, Element g2, Element g2_e_alpha, String file_cipher, String sp) {
        /**            分词，提取关键字           **/
        List<String> keywordList = HanLP.extractKeyword(readFileContent(sp), l1);    //W ={w1, w2, ..., wl1 }
        String ss = String.join(",", keywordList);
        System.out.println(keywordList);
        System.out.println(ss);

        /**            a) With the keywords set W, calculate H'(w1), H'(w2), ..., H'(wl1)           **/
        List<BigInteger> bHw_list = new ArrayList<BigInteger>();
        for (Iterator<String> it = keywordList.iterator(); it.hasNext(); ) {
            byte[] tmp = hashFromStringToZp(it.next(), bp).toBytes();  //H'(w1), H'(w2), ..., H'(wl1)
            String new_Hw = hexBytesToString(tmp);
            BigInteger bigInt = new BigInteger(new_Hw, 16);

            bHw_list.add(bigInt);   //H'(w1), H'(w2), ..., H'(wl1)
        }

        /**            r(x) = γl1x^l1+γl1−1x^(l1−1)+... + γ0，求多项式系数（假设l1为5）**/
        List<BigInteger> bγ_list = calculate_r5(bHw_list);

        /**            b) Randomly selects a random nonce ρ, computes the two parts ciphertexts C1 =e(g1, h)^ρ = e(g1^ρ, h), C2 = g2^ρ            **/
        Element ρ = Zq.newRandomElement().getImmutable();
        BigInteger bρ = ρ.duplicate().toBigInteger();

        Element g1_e_ρ = g1.duplicate().powZn(ρ.duplicate());
        Element C1_prime = bp.pairing(g1_e_ρ.duplicate(), g2_e_alpha.duplicate());   //h = g2_e_alpha
        Element C2_prime = g2.duplicate().powZn(ρ.duplicate());
        String C1 = hexBytesToString(C1_prime.duplicate().toBytes());  //CT---C1
        String C2 = hexBytesToString(C2_prime.duplicate().toBytes());  //CT---C2

        /**  Ci = g2^(γρ)  CT = {C1, C2, Ci(i = 0, ..., l1), CF }  **/
        String CT = C1 + "," + C2 + ",";
//        List<String> Ci_list = new ArrayList<String>();
        for (int i = 0; i <= l1; i++) {
            Element C = g2.duplicate().powZn(ρ.duplicate().mul(bγ_list.get(i)));
            CT = CT + hexBytesToString(C.duplicate().toBytes()) + ",";
//            Ci_list.add(hexBytesToString(C.duplicate().toBytes()));   //CT---C0……Cl1
        }

        CT = CT + file_cipher;
        System.out.println(CT);

        Map<String, String> map = new HashMap<String, String>();
        map.put("keywords", ss);
        map.put("CT", CT);

        return map;
    }

    public static List<BigInteger> calculate_r5(List<BigInteger> bHw_list) {
        BigInteger x1 = bHw_list.get(0);
        BigInteger x2 = bHw_list.get(1);
        BigInteger x3 = bHw_list.get(2);
        BigInteger x4 = bHw_list.get(3);
        BigInteger x5 = bHw_list.get(4);

        BigInteger γ5 = new BigInteger("1");
        BigInteger γ4 = new BigInteger("1");
        BigInteger γ3 = new BigInteger("1");
        BigInteger γ2 = new BigInteger("1");
        BigInteger γ1 = new BigInteger("1");
        BigInteger γ0 = new BigInteger("1");

        γ4 = x1.add(x2.add(x3.add(x4.add(x5)))).negate();
        γ3 = x1.multiply(x2).add(x1.multiply(x3).add(x1.multiply(x4).add(x1.multiply(x5).add(x2.multiply(x3).add(x2.multiply(x4).add(x2.multiply(x5).add(x3.multiply(x4).add(x3.multiply(x5).add(x4.multiply(x5))))))))));
        γ2 = x1.multiply(x2.multiply(x3)).add(x1.multiply(x3.multiply(x4)).add(x1.multiply(x4.multiply(x5)).add(x1.multiply(x5.multiply(x3)).add(x2.multiply(x3.multiply(x4)).add(x2.multiply(x4.multiply(x5)).add(x2.multiply(x5.multiply(x3)).add(x3.multiply(x4.multiply(x5)).add(x1.multiply(x2.multiply(x4)).add(x1.multiply(x2.multiply(x5))))))))))).negate();
        γ1 = x1.multiply(x2.multiply(x3.multiply(x4))).add(x1.multiply(x2.multiply(x3.multiply(x5))).add(x1.multiply(x2.multiply(x4.multiply(x5))).add(x1.multiply(x3.multiply(x4.multiply(x5))).add(x2.multiply(x3.multiply(x4.multiply(x5)))))));
        γ0 = x1.multiply(x2.multiply(x3.multiply(x4.multiply(x5)))).subtract(new BigInteger("1")).negate();

        List<BigInteger> bγ_list = new ArrayList<BigInteger>();
        bγ_list.add(γ0);
        bγ_list.add(γ1);
        bγ_list.add(γ2);
        bγ_list.add(γ3);
        bγ_list.add(γ4);
        bγ_list.add(γ5);

        return bγ_list;
    }

    public static Map<String, String> file_encryption(Field GT, Pairing bp, String sp) throws NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, NoSuchProviderException, InvalidKeyException {
        //selects a random num , computes secret key
        Element T = GT.newRandomElement().getImmutable();
        Element kF = hashFromGTToZp(T.duplicate(), bp);  // secret key kF = h(τ )

        byte[] kf_bytes = kF.toBytes();
        String filecontext = "";
        String file_cipher = "";

        if (SM4algorithm) {
            byte[] bs = new byte[SM4_BLOCK_SIZE];
            System.arraycopy(kf_bytes, 0, bs, 0, SM4_BLOCK_SIZE);
            filecontext = readFileContent(sp);

            sm4_cipherText = SM4Util.encrypt_ECB_Padding(bs, filecontext.getBytes());
            file_cipher = hexBytesToString(sm4_cipherText);   //CF = Enc kF(file).   CT---CF
        }

        Map<String, String> map = new HashMap<String, String>();
        map.put("ciph", file_cipher);
        map.put("KeyF", hexBytesToString(kf_bytes));

        return map;
    }


    //h : GT → Zq*
    public static Element hashFromGTToZp(Element gt_element, Pairing pairing) {
        // h(y) : GT -> Zq
        byte[] gt_bytes = gt_element.getImmutable().toCanonicalRepresentation();
        byte[] zq_bytes = gt_bytes;
        try {
            MessageDigest hasher = MessageDigest.getInstance("SHA-512");
            zq_bytes = hasher.digest(gt_bytes);   //先把GT元素hash成512bits
        } catch (Exception e) {
            e.printStackTrace();
        }
        //再把hash后的bits映射到Zp
        Element hash_result = pairing.getZr().newElementFromHash(zq_bytes, 0, zq_bytes.length).getImmutable();
        return hash_result;
    }

    //H': {0,1}*→ Zq*
    public static Element hashFromStringToZp(String str, Pairing pairing) {
        return pairing.getZr().newElement().setFromHash(str.getBytes(), 0, str.length()).getImmutable();
    }

    public static Element hashFromBytesToZp(byte[] bytes, Pairing pairing) {
        return pairing.getZr().newElement().setFromHash(bytes, 0, bytes.length).getImmutable();
    }

    //16进制的byte[]数组转换为字符串
    public static String hexBytesToString(byte[] bytes) {
        char[] hexChars = new char[bytes.length * 2];
        for (int j = 0; j < bytes.length; j++) {
            int v = bytes[j] & 0xFF;
            hexChars[j * 2] = HEX_ARRAY[v >>> 4];
            hexChars[j * 2 + 1] = HEX_ARRAY[v & 0x0F];
        }
        return new String(hexChars);
    }

    //16进制的字符串转换为byte[]数组
    public static byte[] hexStringToBytes(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                    + Character.digit(s.charAt(i + 1), 16));
        }
        return data;
    }

    public static String readFileContent(String fileName) {
        File file = new File(fileName);
        BufferedReader reader = null;
        StringBuffer sbf = new StringBuffer();
        try {
            reader = new BufferedReader(new FileReader(file));
            String tempStr;
            while ((tempStr = reader.readLine()) != null) {
                sbf.append(tempStr);
                sbf.append("\r\n");
            }
            reader.close();
            return sbf.toString();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
        return sbf.toString();
    }

    public static String writeFileContent(String fileName, String text) {
        try {
            BufferedWriter out = new BufferedWriter(new FileWriter(fileName));
            out.write(text);
            out.close();
            System.out.println("文件创建成功！");
        } catch (IOException e) {
        }
        return "";
    }


}
